﻿using System;
using System.Collections.Generic;
using System.Text;

namespace MLForgeSharp.Models.SupervisedLearningModels.DecisionTree
{
    /// <summary>
    /// ID3模型
    /// </summary>
    public class ID3Model
    {
        private class TreeNode
        {
            public string FeatureName { get; set; } // 分裂特征
            public Dictionary<string, TreeNode> Children { get; set; } // 子节点
            public string Label { get; set; } // 叶节点的标签

            public TreeNode()
            {
                Children = new Dictionary<string, TreeNode>();
            }
        }

        private TreeNode root; // 决策树的根节点

        public ID3Model()
        {
            root = new TreeNode();
        }

        // 计算熵
        private double CalculateEntropy(int[] counts)
        {
            double entropy = 0.0;
            int total = 0;
            foreach (int count in counts)
            {
                total += count;
            }
            foreach (int count in counts)
            {
                if (count > 0)
                {
                    double probability = (double)count / total;
                    entropy -= probability * Math.Log(probability, 2);
                }
            }
            return entropy;
        }

        // 计算信息增益
        private double CalculateInformationGain(List<string[]> data, int featureIndex, int labelIndex)
        {
            Dictionary<string, List<string[]>> subsets = new Dictionary<string, List<string[]>>();
            foreach (string[] row in data)
            {
                string featureValue = row[featureIndex];
                if (!subsets.ContainsKey(featureValue))
                {
                    subsets[featureValue] = new List<string[]>();
                }
                subsets[featureValue].Add(row);
            }

            double totalEntropy = CalculateEntropy(GetLabelCounts(data, labelIndex));
            double splitEntropy = 0.0;

            foreach (var subset in subsets.Values)
            {
                double subsetProbability = (double)subset.Count / data.Count;
                splitEntropy += subsetProbability * CalculateEntropy(GetLabelCounts(subset, labelIndex));
            }

            return totalEntropy - splitEntropy;
        }

        // 获取标签的计数
        private int[] GetLabelCounts(List<string[]> data, int labelIndex)
        {
            Dictionary<string, int> labelCounts = new Dictionary<string, int>();
            foreach (string[] row in data)
            {
                string label = row[labelIndex];
                if (!labelCounts.ContainsKey(label))
                {
                    labelCounts[label] = 0;
                }
                labelCounts[label]++;
            }

            int[] counts = new int[labelCounts.Count];
            int i = 0;
            foreach (int count in labelCounts.Values)
            {
                counts[i++] = count;
            }
            return counts;
        }

        // 构建决策树
        private TreeNode BuildTree(List<string[]> data, List<string> featureNames, int labelIndex)
        {
            if (data.Count == 0)
            {
                return new TreeNode { Label = "Unknown" };
            }

            int[] labelCounts = GetLabelCounts(data, labelIndex);
            if (labelCounts.Length == 1)
            {
                return new TreeNode { Label = data[0][labelIndex] };
            }

            if (featureNames.Count == 0)
            {
                return new TreeNode { Label = GetMajorityLabel(data, labelIndex) };
            }

            int bestFeatureIndex = 0;
            double bestInformationGain = 0.0;
            for (int i = 0; i < featureNames.Count; i++)
            {
                double informationGain = CalculateInformationGain(data, i, labelIndex);
                if (informationGain > bestInformationGain)
                {
                    bestInformationGain = informationGain;
                    bestFeatureIndex = i;
                }
            }

            TreeNode node = new TreeNode { FeatureName = featureNames[bestFeatureIndex] };
            Dictionary<string, List<string[]>> subsets = SplitData(data, bestFeatureIndex);

            List<string> remainingFeatures = new List<string>(featureNames);
            remainingFeatures.RemoveAt(bestFeatureIndex);

            foreach (var subset in subsets)
            {
                node.Children[subset.Key] = BuildTree(subset.Value, remainingFeatures, labelIndex);
            }

            return node;
        }

        // 获取多数标签
        private string GetMajorityLabel(List<string[]> data, int labelIndex)
        {
            Dictionary<string, int> labelCounts = new Dictionary<string, int>();
            foreach (string[] row in data)
            {
                string label = row[labelIndex];
                if (!labelCounts.ContainsKey(label))
                {
                    labelCounts[label] = 0;
                }
                labelCounts[label]++;
            }

            string majorityLabel = "";
            int maxCount = 0;
            foreach (var kvp in labelCounts)
            {
                if (kvp.Value > maxCount)
                {
                    maxCount = kvp.Value;
                    majorityLabel = kvp.Key;
                }
            }
            return majorityLabel;
        }

        // 根据特征值分割数据
        private Dictionary<string, List<string[]>> SplitData(List<string[]> data, int featureIndex)
        {
            Dictionary<string, List<string[]>> subsets = new Dictionary<string, List<string[]>>();
            foreach (string[] row in data)
            {
                string featureValue = row[featureIndex];
                if (!subsets.ContainsKey(featureValue))
                {
                    subsets[featureValue] = new List<string[]>();
                }
                subsets[featureValue].Add(row);
            }
            return subsets;
        }

        // 训练模型
        public void Train(List<string[]> data, List<string> featureNames, int labelIndex)
        {
            root = BuildTree(data, featureNames, labelIndex);
        }

        // 预测
        public string Predict(string[] features, List<string> featureNames)
        {
            TreeNode node = root;
            while (node.Label == null)
            {
                string featureValue = features[featureNames.IndexOf(node.FeatureName)];
                if (!node.Children.ContainsKey(featureValue))
                {
                    return "Unknown";
                }
                node = node.Children[featureValue];
            }
            return node.Label;
        }
    }

    // 示例程序
    public class ID3ModelExample
    {
        public ID3ModelExample()
        {
            // 示例数据
            List<string[]> data = new List<string[]>
        {
            new string[] { "Sunny", "Hot", "High", "Weak", "No" },
            new string[] { "Sunny", "Hot", "High", "Strong", "No" },
            new string[] { "Overcast", "Hot", "High", "Weak", "Yes" },
            new string[] { "Rain", "Mild", "High", "Weak", "Yes" },
            new string[] { "Rain", "Cool", "Normal", "Weak", "Yes" },
            new string[] { "Rain", "Cool", "Normal", "Strong", "No" },
            new string[] { "Overcast", "Cool", "Normal", "Strong", "Yes" },
            new string[] { "Sunny", "Mild", "High", "Weak", "No" },
            new string[] { "Sunny", "Cool", "Normal", "Weak", "Yes" },
            new string[] { "Rain", "Mild", "Normal", "Weak", "Yes" },
            new string[] { "Sunny", "Mild", "Normal", "Strong", "Yes" },
            new string[] { "Overcast", "Mild", "High", "Strong", "Yes" },
            new string[] { "Overcast", "Hot", "Normal", "Weak", "Yes" },
            new string[] { "Rain", "Mild", "High", "Strong", "No" }
        };

            List<string> featureNames = new List<string> { "Outlook", "Temperature", "Humidity", "Wind" };
            int labelIndex = 4; // 标签列索引

            // 创建模型
            ID3Model model = new ID3Model();

            // 训练模型
            model.Train(data, featureNames, labelIndex);

            // 预测
            string[] testFeatures = { "Sunny", "Mild", "High", "Strong" };
            string predictedLabel = model.Predict(testFeatures, featureNames);

            System.Console.WriteLine("预测标签: " + predictedLabel);
        }
    }
}
